//
//  trie.swift
//  others
//
//  Created by lhc on 2024/8/17.
//

import Foundation


class Trie {
    
    class Node {
        var isWord = false
        var map = [Character: Node]()
    }
    
    private var root = Node()
    private var size = 0
    
    
    func getSize() -> Int {
        return self.size
    }
    
    func add(str: String) {
        
        var cur = self.root
        for char in str {
            // 如果当前结点的字典里不包括当前char 那么创建一个新结点
            if !cur.map.keys.contains(char) {
                cur.map[char] = Node()
            }
            cur = cur.map[char]!
        }
        
        // 遍历结束后说明添加新词结束,这时候如果当前结点的isWord不是True，说明添加了一个新词
        if !cur.isWord {
            cur.isWord = true
            self.size += 1
        }
    }
    
    func contain(str: String) -> Bool {
        
        var cur = self.root
        for char in str {
            
            if !cur.map.keys.contains(char) {
                return false
            }
            cur = cur.map[char]!
        }
        
        return cur.isWord
    }
    
    
    func isPrefix(_ prefix: String) -> Bool {
        
        var cur = self.root
        for char in prefix {
            if !cur.map.keys.contains(char) {
                return false
            }
            cur = cur.map[char]!
        }
        
        return true
    }
}
